import streamlit as st
import tensorflow as tf
import numpy as np
from PIL import Image
import cv2

# 设置页面标题
st.set_page_config(page_title="垃圾分类识别系统", layout="wide")

# 加载模型
@st.cache_resource
def load_model():
    try:
        model = tf.keras.models.load_model('./garbage_classifier.h5')
        return model
    except:
        st.error('请先训练模型并生成 garbage_classifier.h5 文件')
        return None

# 图片预处理
def preprocess_image(image):
    # 调整图片大小
    image = image.resize((224, 224))
    # 转换为numpy数组
    image_array = np.array(image)
    # 确保图片是RGB格式
    if image_array.shape[-1] == 4:
        image_array = image_array[:,:,:3]
    # 扩展维度
    image_array = np.expand_dims(image_array, axis=0)
    # 归一化
    image_array = image_array / 255.0
    return image_array

# 主函数
def main():
    st.title('垃圾分类识别系统')
    st.write('上传图片，系统将自动识别垃圾类别')

    # 创建两列布局
    col1, col2 = st.columns(2)

    # 文件上传
    uploaded_file = col1.file_uploader("选择图片文件", type=['jpg', 'jpeg', 'png'])

    if uploaded_file is not None:
        # 显示上传的图片
        image = Image.open(uploaded_file)
        col1.image(image, caption='上传的图片', use_column_width=True)

        # 加载模型
        model = load_model()

        if model is not None:
            # 预处理图片
            processed_image = preprocess_image(image)

            # 预测
            prediction = model.predict(processed_image)
            classes = ['可回收垃圾', '有害垃圾', '厨余垃圾', '其他垃圾']
            predicted_class = classes[np.argmax(prediction)]
            confidence = np.max(prediction) * 100

            # 显示预测结果
            col2.subheader('预测结果')
            col2.write(f'类别: {predicted_class}')
            col2.write(f'置信度: {confidence:.2f}%')

            # 显示所有类别的概率
            col2.subheader('各类别概率')
            for i, (class_name, prob) in enumerate(zip(classes, prediction[0])):
                col2.progress(float(prob))
                col2.write(f'{class_name}: {prob*100:.2f}%')

if __name__ == '__main__':
    main()